import os
import re
import subprocess
import signal
from typing import List, Tuple

DESIGN_NAME="arithmetic_core_double_fpu_verilog"
DESIGN_TOP="fpu"
CLK_PORTS="clk"

PROJECT_DIR = os.getenv("PROJECT_DIR")

DESIGN_DIR = f"{PROJECT_DIR}/{DESIGN_NAME}"
DESIGN_PATHS=" \\\n    ".join([os.path.join(DESIGN_DIR,"rtl", p) for p in os.listdir(os.path.join(DESIGN_DIR,"rtl")) if p.endswith(".v")])

CLK_PERIOD=1
abc_scripts = [
    "strash;balance;rewrite;refactor;balance;rewrite,-z;refactor,-z;balance;map,-v",
    "strash;resub;resub,-z;rewrite,-z;balance;resub,-z;balance;map,-v",
    "strash;dc2;balance;rewrite;balance;dc2;balance;map,-v",
    "strash;&get,-n;&fraig,-x;&put;resub;resub,-z;rewrite;balance;map,-v",
    "strash;rewrite;refactor;balance;rewrite,-z;dc2;balance;rewrite;map,-v",
    "strash;rewrite;rewrite;refactor;refactor,-z;rewrite,-z;balance;map,-v",
    "strash;dc2;balance;dretime;rewrite;balance;dc2;map,-v",
    "strash;scorr;rewrite;refactor;balance;dch;balance;map,-v",
    "strash;dc2;rewrite,-z;resub,-z;balance;refactor,-z;balance;map,-v",
    "strash;&get,-n;&fraig,-x;&put;scorr;dc2;dretime;strash;&get,-n;&dch,-f;&nf;&put;map,-v",
]
yosys_lib_imports = {
    "nangate45": f"-liberty {PROJECT_DIR}/liberty/nangate45/lib/NangateOpenCellLibrary_typical.lib",
    # "skywater130": f"-liberty {PROJECT_DIR}/liberty/skywater130/lib/sky130_fd_sc_hd__tt_025C_1v80.lib",
}
opensta_lib_imports = {
    "nangate45": f"read_liberty {PROJECT_DIR}/liberty/nangate45/lib/NangateOpenCellLibrary_typical.lib",
    # "skywater130": f"read_liberty {PROJECT_DIR}/liberty/skywater130/lib/sky130_fd_sc_hd__tt_025C_1v80.lib",
}



def run_with_timeout(cmd: List[str], timeout: int) -> Tuple[int, str, str]:
    """
    运行命令并在超时时杀掉整个进程组。

    参数:
        cmd: 命令和参数的列表，例如 ["iverilog", "test.v"]
        timeout: 超时时间（秒）

    返回:
        (returncode, stdout, stderr)
    """
    try:
        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            start_new_session=True
        )
        stdout, stderr = proc.communicate(timeout=timeout)
        return proc.returncode, stdout, stderr
    except subprocess.TimeoutExpired:
        os.killpg(proc.pid, signal.SIGTERM)
        try:
            stdout, stderr = proc.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            os.killpg(proc.pid, signal.SIGKILL)
            stdout, stderr = proc.communicate()
        return -9, stdout, stderr

def parse_area_report(rpt_path: str) -> float:
    with open(rpt_path, "r") as f:
        report = f.read()
    match = re.search(rf"^\s*Chip area for module '\\{DESIGN_TOP}':\s+(\d+\.\d+)", report, re.MULTILINE)
    if match:
        return float(match.group(1))
    else:
        print("Could not find total cell area in area report")
        return None

def parse_delay_report(rpt_path: str) -> float:
    with open(rpt_path, "r") as f:
        report = f.read()
    match = re.search(r"^\s+(\d+\.\d+)\s+data arrival time", report, re.MULTILINE)
    if match:
        return float(match.group(1))
    else:
        print("Could not find data arrival time in delay report")
        return None

def parse_power_report(rpt_path: str) -> float:
    with open(rpt_path, "r") as f:
        report = f.read()
    match = re.search(r"^\s*Total\s+\d+\.\d+(?:e-\d+)?\s+\d+\.\d+(?:e-\d+)?\s+\d+\.\d+(?:e-\d+)?\s+(\d+\.\d+(?:e-\d+)?\s+)", report, re.MULTILINE)
    if match:
        return float(match.group(1))
    else:
        print("Could not find total power in power report")
        return None

def synthesis():
    os.makedirs(f"{DESIGN_DIR}/scripts/", exist_ok=True)
    os.makedirs(f"{DESIGN_DIR}/syn/", exist_ok=True)
    os.makedirs(f"{DESIGN_DIR}/rpt/", exist_ok=True)
    for lib_name in yosys_lib_imports.keys():
        os.makedirs(f"{DESIGN_DIR}/scripts/{lib_name}/", exist_ok=True)
        os.makedirs(f"{DESIGN_DIR}/syn/{lib_name}", exist_ok=True)
        os.makedirs(f"{DESIGN_DIR}/rpt/{lib_name}", exist_ok=True)
        for i, ABC_SCRIPT in enumerate(abc_scripts):
            NETLIST_PATH=f"{DESIGN_DIR}/syn/{lib_name}/netlist_{i}.v"
            AREA_RPT_PATH=f"{DESIGN_DIR}/rpt/{lib_name}/area_{i}.rpt"
            DELAY_RPT_PATH=f"{DESIGN_DIR}/rpt/{lib_name}/delay_{i}.rpt"
            POWER_RPT_PATH=f"{DESIGN_DIR}/rpt/{lib_name}/power_{i}.rpt"
            with open(f"{PROJECT_DIR}/eda_script_templates/synth.ys", "r") as f:
                template = f.read()
            yosys_script = template.format(
                DESIGN_PATHS=DESIGN_PATHS,
                NETLIST_PATH=NETLIST_PATH,
                YOSYS_LIB_IMPORTS=yosys_lib_imports[lib_name],
                ABC_SCRIPT=ABC_SCRIPT,
                AREA_RPT_PATH=AREA_RPT_PATH,
            )
            with open(f"{DESIGN_DIR}/scripts/{lib_name}/synth_{i}.ys", "w") as f:
                f.write(yosys_script)
            yosys_cmd = ["yosys", "-q", "-s", f"{DESIGN_DIR}/scripts/{lib_name}/synth_{i}.ys"]
            retcode, stdout, stderr = run_with_timeout(yosys_cmd, timeout=300)
            if retcode != 0:
                print(f"Yosys synthesis failed for {lib_name} script {i}")
                print("STDOUT:", stdout)
                print("STDERR:", stderr)
                area = None
                continue
            area = parse_area_report(AREA_RPT_PATH)
            with open(f"{PROJECT_DIR}/eda_script_templates/delay.tcl", "r") as f:
                template = f.read()
            delay_script = template.format(
                OPENSTA_LIB_IMPORTS=opensta_lib_imports[lib_name],
                NETLIST_PATH=NETLIST_PATH,
                DESIGN_TOP=DESIGN_TOP,
                CLK_PORTS=CLK_PORTS,
                DELAY_RPT_PATH=DELAY_RPT_PATH,
            )
            with open(f"{DESIGN_DIR}/scripts/{lib_name}/delay_{i}.tcl", "w") as f:
                f.write(delay_script)
            sta_cmd = ["sta", "-exit", f"{DESIGN_DIR}/scripts/{lib_name}/delay_{i}.tcl"]
            retcode, stdout, stderr = run_with_timeout(sta_cmd, timeout=120)
            if retcode != 0:
                print(f"OpenSTA timing analysis failed for {lib_name} script {i}")
                print("STDOUT:", stdout)
                print("STDERR:", stderr)
                delay = None
                continue
            delay = parse_delay_report(DELAY_RPT_PATH)
            CLK_PERIOD = delay
            with open(f"{PROJECT_DIR}/eda_script_templates/power.tcl", "r") as f:
                template = f.read()
            power_script = template.format(
                READ_LIBERTY=opensta_lib_imports[lib_name],
                NETLIST_PATH=NETLIST_PATH,
                DESIGN_TOP=DESIGN_TOP,
                CLK_PERIOD=CLK_PERIOD,
                CLK_PORTS=CLK_PORTS,
                POWER_RPT_PATH=POWER_RPT_PATH,
            )
            with open(f"{DESIGN_DIR}/scripts/{lib_name}/power_{i}.tcl", "w") as f:
                f.write(power_script)
            power_cmd = ["sta", "-exit", f"{DESIGN_DIR}/scripts/{lib_name}/power_{i}.tcl"]
            retcode, stdout, stderr = run_with_timeout(power_cmd, timeout=120)
            if retcode != 0:
                print(f"OpenSTA power analysis failed for {lib_name} script {i}")
                print("STDOUT:", stdout)
                print("STDERR:", stderr)
                power = None
                continue
            power = parse_power_report(POWER_RPT_PATH)
            with open(f"{DESIGN_DIR}/rpt/{lib_name}/summary.csv", "a") as f:
                if i == 0:
                    f.write("Area,Delay,Power\n")
                f.write(f"{area},{delay},{power}\n")
            print(f"Synthesis and analysis completed for {lib_name} script {i}")
            

if __name__ == "__main__":
    synthesis()